import matplotlib.axes as maxes
#from matplotlib.path import Path
from matplotlib.transforms import Transform
from matplotlib.cbook import is_string_like

import numpy as np
import pywcsgrid2 as pywcsgrid
from pywcsgrid2.coords import coord_conv

from  axislines import AxisLineHelper, GridHelperRectlinear, AxisLine

from parasite_axes import HostAxes, ParasiteAxesAuxTrans
from pywcsgrid2.coords.coord_system import coord_system_guess, coord_sys, FK5, FK4, GAL


import re

_wcs_key_pattern = re.compile(r'^(NAXIS|CD|CDELT|CRPIX|CRVAL|CTYPE|CROTA|LONGPOLE|LATPOLE|PV|DISTORT|OBJECT|BUNIT|EPOCH|EQUINOX|LTV|LTM|DTV|DTM)')

def _filterwcs(h):
    """
    select wcs related cards
    """
    l = [card for card in h.ascardlist() if _wcs_key_pattern.match(card.key)]
    return pywcs.WCS(pyfits.Header(l))


_coord_sys_dict = dict(fk5=FK5,
                       fk4=FK4,
                       gal=GAL)

class _pgtext_to_latex(object):
    p = re.compile(r"\\u(\w*)\\d")
    _convert_table = {"h":r"\mathrm{h}",
                      "m":r"\mathrm{m}",
                      "s":r"\mathrm{s}",
                      "o":r"{\circ}",
                      "'":r"{\prime}",
                      '"':r"{\prime\prime}",
                      }

    def __call__(self, s):
        startindex=0
        r = []
        if not s: return ""
        for m in self.p.finditer(s):
            r.append(s[startindex:m.start()])
            c = m.groups()[0]
            r.append("^{" + self._convert_table.get(c, c) + "}\,")
            startindex=m.end()
        r.append(s[startindex:])
        r = "".join(r)
        r = r.replace("'", r"^{\prime}").replace('"',r"^{\prime\prime}")

        return "$"+r+"$"

_pgtext_to_latex = _pgtext_to_latex()


class FixedAxisLineHelper(AxisLineHelper.Fixed):

    def __init__(self, axes, side, nth_coord, nth_coord_ticks):
        """
        nth_coord = along which coordinate value varies
        in 2d, nth_coord = 0 ->  x axis, nth_coord = 1 -> y axis
        """

        AxisLineHelper.Fixed.__init__(self, axes, loc=side,
                                      nth_coord=nth_coord,
                                      passingthrough_point=None,
                                      label_direction=None)

        #self.nth_coord = nth_coord
        if nth_coord_ticks is None:
            nth_coord_ticks = self.nth_coord
        self.nth_coord_ticks = nth_coord_ticks

        self.side = side


    def get_tick_transform(self):
        return self.axes.transData

    def get_tick_iterators(self):
        """tick_loc, tick_angle, tick_label"""

        g = self.axes._grid_helper

        axis = [g.wcsgrid.axis1, g.wcsgrid.axis2][self.nth_coord_ticks]
        ticks1 = axis["ticks_"+self.side]
        tickangles1 = axis["tickangles_"+self.side]
        ticklabels1 = axis["ticklabels_"+self.side]

        transform = self.axes.transData
        def f1():
            for (cx, cy), a, l in zip(ticks1, tickangles1, ticklabels1):
                l = _pgtext_to_latex(l)
                if a:
                    yield [cx, cy], a/3.14159*180.+90, l

        axis = [g.wcsgrid.axis1, g.wcsgrid.axis2][1-self.nth_coord_ticks]
        ticks2 = axis["ticks_"+self.side]
        tickangles2 = axis["tickangles_"+self.side]
        ticklabels2 = axis["ticklabels_"+self.side]

        def f2():
            for (cx, cy), a, l in zip(ticks2, tickangles2, ticklabels2):
                l = _pgtext_to_latex(l)
                if a:
                    yield [cx, cy], a/3.14159*180.+90, ""

        return f1(), f2()


class GridHelperWcs(GridHelperRectlinear):

    def __init__(self, wcs):
        super(GridHelperWcs, self).__init__()

        self._wcs = wcs
        self.wcsgrid = None
        self._old_values = None
        self._wcsgrid_params = dict(coord_format="GE",
                                    label_density=(6, 6),
                                    grid1=[], grid2=[])

        coord_guess = coord_system_guess(wcs.ctype[0],
                                         wcs.ctype[1],
                                         wcs.equinox)
        if coord_guess is None:
            raise ValueError("")
        else:
            self._wcsgrid_orig_coord_system = coord_guess
        #self._wcsgrid_orig_coord_system = None
        self._wcsgrid_display_coord_system = None


    def update_wcsgrid_params(self, **ka):
        """
        coord_format="GE",
        label_density=(6, 6),
        grid1=[], grid2=[]):
        """

        for k in ["coord_format", "label_density", "grid1", "grid2", "src_system", "dest_system"]:
            if k in ka:
                self._wcsgrid_params[k] = ka.pop(k)

        if len(ka) > 0:
            raise ValueError("keyword name should be one of coord_format, label_densit, grid1, grid2")


    def get_wcsgrid_params(self):
        return self._wcsgrid_params

    def _update(self, x1, x2, y1, y2):

        # update wcsgrid

        if self._force_update is False and \
               self._old_values == (x1, x2, y1, y2,
                                    self.get_wcsgrid_params()):
            return

        if self._wcsgrid_display_coord_system is not None:
            src_system = self._wcsgrid_orig_coord_system
            tgt_system = self._wcsgrid_display_coord_system
        else:
            src_system = None
            tgt_system = None

        self.update_wcsgrid_params(src_system=src_system,
                                   dest_system=tgt_system)


        self.wcsgrid = pywcsgrid.grid(self._wcs, x1, y1, x2, y2,
                                      **self._wcsgrid_params)

        self._old_values = (x1, x2, y1, y2, self.get_wcsgrid_params().copy())

        self._force_update = False


    def get_gridlines(self):
        return self.wcsgrid.grid_lines


    def get_axisline_helper(self, side, nth_coord, nth_coord_ticks=None):
        return FixedAxisLineHelper(self.axes, side, nth_coord, nth_coord_ticks)

    def get_axisline(self, loc,
                     nth_coord=None, passthrough_point=None,
                     transform=None,
                     tick_direction="in",
                     label_direction=None,
                     offset_transform=None):

        _helper = self.get_axisline_helper(loc, nth_coord)

        axisline = AxisLine(self.axes, _helper)

        return axisline


    def set_display_coord_system(self, coord_system):
        if self._wcsgrid_orig_coord_system is None:
            raise ValueError("_pywcsgrid_orig_coord_system is not set")

        elif is_string_like(coord_system):
            coord_system = _coord_sys_dict[coord_system.lower()]

        self._wcsgrid_display_coord_system = coord_system
        self.invalidate()


from matplotlib.path import Path

class CurvedTransform(Transform):
    def __init__(self, resolution=None):
        """
        Create a new WCS transform.
        """
        Transform.__init__(self)
        if resolution is None:
            resolution = 100.
        self._resolution = resolution

    def transform_path(self, path):
        vertices = path.vertices
        ipath = path.interpolated(self._resolution)
        return Path(self.transform(ipath.vertices), ipath.codes)

    transform_path.__doc__ = Transform.transform_path.__doc__

    transform_path_non_affine = transform_path
    transform_path_non_affine.__doc__ = Transform.transform_path_non_affine.__doc__

class WcsSky2PixelTransform(CurvedTransform):
    """
    """
    input_dims = 2
    output_dims = 2
    is_separable = False

    def __init__(self, wcs, resolution=None):
        """
        Create a new WCS transform.
        """
        CurvedTransform.__init__(self, resolution)

        if wcs.naxis > 2:
            wcs = wcs.wcssub([1, 2])

        self.wcs = wcs

    def transform(self, ll):
        return self.wcs.world2pixel(ll)


    transform.__doc__ = Transform.transform.__doc__

    transform_non_affine = transform
    transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__


    def inverted(self):
        return WcsPixel2SkyTransform(self.wcs, self._resolution)
    inverted.__doc__ = Transform.inverted.__doc__





class WcsPixel2SkyTransform(CurvedTransform):
    """
    The base Aitoff transform.
    """
    input_dims = 2
    output_dims = 2
    is_separable = False

    def __init__(self, wcs, resolution=None):
        """
        Create a new WCS transform.  Resolution is the number of steps
        to interpolate between each input line segment to approximate its
        path.
        """
        CurvedTransform.__init__(self, resolution)

        if wcs.naxis > 2:
            wcs = wcs.wcssub([1, 2])

        self.wcs = wcs

    def transform(self, ll):
        return self.wcs.pixel2world(ll)


    transform.__doc__ = Transform.transform.__doc__

    transform_non_affine = transform
    transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__


    def inverted(self):
        return WcsSky2PixelTransform(self.wcs, self._resolution)
    inverted.__doc__ = Transform.inverted.__doc__



class WcsSky2SkyTransform(CurvedTransform):
    """
    The base Aitoff transform.
    """
    input_dims = 2
    output_dims = 2
    is_separable = False

    def __init__(self, src_coord, dest_coord, resolution=None):
        """
        Create a new WCS transform.  Resolution is the number of steps
        to interpolate between each input line segment to approximate its
        path.
        """

        CurvedTransform.__init__(self, resolution)

        self.src_coord = src_coord
        self.dest_coord = dest_coord

        self.coord_conv_func = coord_conv(src_coord, dest_coord)


    def transform(self, ll):
        lat1, lon1 = ll[:,0], ll[:,1]
        lat2, lon2 = self.coord_conv_func(lat1, lon1)

        return np.concatenate((lat2[:,np.newaxis], lon2[:,np.newaxis]),
                              1)


    transform.__doc__ = Transform.transform.__doc__

    transform_non_affine = transform
    transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__


    def inverted(self):
        return WcsSky2SkyTransform(self.dest_coord, self.src_coord)
    inverted.__doc__ = Transform.inverted.__doc__



class AxesWcs(HostAxes):
    def __init__(self, *kl, **kw):

        if "grid_helper" not in kw:
            self._wcs = kw.pop("wcs", None)
            if self._wcs is None:
                raise ValueError("wcs")

            grid_helper = GridHelperWcs(self._wcs)
            kw["grid_helper"] = grid_helper
        else:
            self._wcs = kw["grid_helper"]._wcs

        super(AxesWcs, self).__init__(*kl, **kw)

        self._init_parasites()



    def _init_parasites(self):
        wcs = self._wcs
        ax = ParasiteAxesAuxTrans(self,
                                  WcsSky2PixelTransform(wcs),
                                  viewlim_mode="equal")
        self._wcsgrid_wcsaxes = {0:ax}
        self.parasites.append(ax)


    def __getitem__(self, key):

        # check if key is a valid coord_sys instance
        if key == 0 or isinstance(key, coord_sys) :
            pass
        elif is_string_like(key):
            key = _coord_sys_dict[key.lower()]
        else:
            raise ValueError("argument key must be an instance of %s" % repr(coord_sys))

        if key not in self._wcsgrid_wcsaxes:
            if self.get_grid_helper()._wcsgrid_orig_coord_system == key:
                self._wcsgrid_wcsaxes[key] = self._wcsgrid_wcsaxes[0]
            else:
                orig_coord = self.get_grid_helper()._wcsgrid_orig_coord_system

                tr = WcsSky2SkyTransform(key, orig_coord)
                ax = ParasiteAxesAuxTrans(self,
                                          tr + self._wcsgrid_wcsaxes[0].transAux)
                self._wcsgrid_wcsaxes[key] = ax


        return self._wcsgrid_wcsaxes[key]



SubplotWcs = maxes.subplot_class_factory(AxesWcs)



if __name__ == "__main__":
    import pyfits, pywcs
    import matplotlib.pyplot as plt
    fig = plt.figure(1)
    fig.clf()
    fname = "/Users/jjlee/local/src/astropy/pywcsgrid_all/test/i013b4h0.fit"
    fname = "/bulk/blackstone1/lee/local/src/astropy/trunk/pywcsgrid_all/test/i013b4h0.fit"
    f = pyfits.open(fname)
    d, h = f[0].data, f[0].header
    d.shape = d.shape[-2:] # this particular image has a shape of [1, 500, 500]
    wcs = pywcs.WCS(h)
    ax1 = SubplotWcs(fig, 1, 2, 1, wcs=wcs)
    ax2 = SubplotWcs(fig, 1, 2, 2, wcs=wcs)
    #grid_helper = GridHelperWcs(wcs)
    #ax = Subplot(fig, 1, 1, 1, grid_helper=grid_helper)

    fig.add_subplot(ax1)
    fig.add_subplot(ax2)

    import matplotlib.colorbar

    im1 = ax1.imshow(d, origin="low", vmin=00, vmax=301,
                     cmap=matplotlib.colorbar.cm.gray_r)

    im2 = ax2.imshow(d, origin="low", vmin=00, vmax=301,
                     cmap=matplotlib.colorbar.cm.gray_r)

    ax1.get_grid_helper().set_display_coord_system("fk5")
    ax2.get_grid_helper().set_display_coord_system("gal")

    ax1.grid()
    ax2.grid()

    ax2.axislines["left"]._helper.nth_coord_ticks=0
    ax2.axislines["right"]._helper.nth_coord_ticks=0
    ax2.axislines["bottom"]._helper.nth_coord_ticks=1
    ax2.axislines["top"]._helper.nth_coord_ticks=1

    plt.draw()

    #print grid.wcsgrid.axis1["grid_lines"]
